import json
import os

from glob import glob, iglob
import re
import base64
from shutil import rmtree
from PIL import Image
import argparse
from numpy import random
import cv2
from os.path import join
import numpy as np


# INPUT_DIR = "/media/jiurui/92466C31466C186D/samples/数码显微镜拍摄化石-崔凤强-20190415/json/"
# OUTPUT_DIR = "/media/jiurui/92466C31466C186D/samples/数码显微镜拍摄化石-崔凤强-20190415/labelme_py/"
# label_DIR = "/media/jiurui/92466C31466C186D/samples/数码显微镜拍摄化石-崔凤强-20190415/label_fill/"
# side_DIR = "/media/jiurui/92466C31466C186D/samples/数码显微镜拍摄化石-崔凤强-20190415/side/"

ALL_dir = r"E:\xianhuang0818\Seg"
INPUT_DIR = os.path.join(ALL_dir, "json")
OUTPUT_DIR = os.path.join(ALL_dir, "labelme_py")
label_DIR = os.path.join(ALL_dir, "masks")
side_DIR = os.path.join(ALL_dir, "side")
side_DIR_INV = os.path.join(ALL_dir, "side_inv")
ROI_DIR = os.path.join(ALL_dir, "roi")
img_DIR = os.path.join(ALL_dir, "images")
GrayEqu_DIR = os.path.join(ALL_dir, "GrayEqu")
dir_voc_img = r"H:\yuanbaoxi\ybx_gitee\pspnet_pytorch\VOCdevkit\xh\JPEGImages"
dir_voc_label = r"H:\yuanbaoxi\ybx_gitee\pspnet_pytorch\VOCdevkit\xh\SegmentationClass"

# INPUT_DIR = r"H:\yuanbaoxi\samples\xianhuang\seg\json"
# OUTPUT_DIR = r"H:\yuanbaoxi\samples\xianhuang\seg\labelme_py"
# label_DIR = r"H:\yuanbaoxi\samples\xianhuang\seg\masks"
# side_DIR = r"H:\yuanbaoxi\samples\xianhuang\seg\side"
# img_DIR = r"H:\yuanbaoxi\samples\xianhuang\seg\images"

def empty_dir(path):
    """ empty specified dir """
    if os.path.exists(path):
        rmtree(path)
    os.mkdir(path)


def ensure_dir(path):
    """ empty specified dir """
    if not os.path.exists(path):
        os.mkdir(path)


def get_bbox(coords):
    """ get bounding box in format [tlx, tly, w, h] """
    min_x = None
    min_y = None
    max_x = None
    max_y = None

    for [x, y] in coords:
        min_x = x if not min_x else min(x, min_x)
        min_y = y if not min_y else min(y, min_y)
        max_x = x if not max_x else max(x, max_x)
        max_y = y if not max_y else max(y, max_y)

    return [min_x, min_y, max_x - min_x, max_y - min_y]



def labelme_2_mask_rcnn_ballon():
    std_size = 473

    # ensure_dir(INPUT_DIR)
    # ensure_dir(OUTPUT_DIR)
    empty_dir(OUTPUT_DIR)
    empty_dir(dir_voc_label)
    empty_dir(dir_voc_img)
    empty_dir(label_DIR)
    empty_dir(side_DIR)
    empty_dir(side_DIR_INV)
    empty_dir(ROI_DIR)
    empty_dir(GrayEqu_DIR)
    """ Browse through all marked json files """
    img_idx = 0
    for file in iglob(INPUT_DIR + '/*.json'):
        with open(file, 'r', encoding='utf-8') as f:

            """ Load json files """
            data = json.load(f)

            """ Save image file """
            j_name = file.split("\\")[-1]
            ts = j_name.split(".")
            title = "%d"%img_idx
            img_idx += 1
            # title = ts[0]
            file_name = join(OUTPUT_DIR, j_name)
            print(file_name, file)
            image_data = base64.b64decode(data["imageData"])
            with open(file_name, 'wb') as fi:
                fi.write(image_data)
            img_ori_bgr = cv2.imread(file_name, cv2.IMREAD_COLOR)
            img_ori = cv2.imread(file_name, cv2.IMREAD_GRAYSCALE)
            img_equ = cv2.equalizeHist(img_ori)
            img_equ = cv2.cvtColor(img_equ, cv2.COLOR_GRAY2BGR)
            save_name = join(GrayEqu_DIR, "%s.jpg" % title)
            img_equ = cv2.resize(img_equ, (std_size, std_size))
            save_name = join(dir_voc_img, title) + ".jpg"
            cv2.imwrite(save_name, img_equ)

            size_ori = img_ori.shape
            """ Get image width x height """
            boxImg = np.zeros((size_ori[0], size_ori[1]), dtype="uint8")
            sideImg = np.zeros((size_ori[0], size_ori[1]), dtype="uint8")
            """ Process each shape (annotation) """
            regions = {}
            reg_id = 0
            for shape in data['shapes']:
                cat = shape['label']
                """ Form segment out of points """
                segment_x = []
                segment_y = []
                for [x, y] in shape['points']:
                    segment_x.append(int(x))
                    segment_y.append(int(y))
                e = []
                for i in range(0,len(segment_x)):
                    e.append([segment_x[i], segment_y[i]])
                e = np.array(e)
                cv2.drawContours(boxImg, [e], -1, (255, 255, 255), cv2.FILLED)
                cv2.drawContours(sideImg, [e], -1, (255, 255, 255), 20)

            save_name = join(label_DIR, title) + ".jpg"
            cv2.imwrite(save_name, boxImg)
            save_name = join(side_DIR, title) + ".png"
            cv2.imwrite(save_name, sideImg)

            save_name = join(dir_voc_label, title) + ".png"
            img_stdlabel = cv2.resize(boxImg, (std_size, std_size))
            cv2.imwrite(save_name, img_stdlabel)
            # img_stdlabel = cv2.resize(sideImg, (std_size, std_size))
            # cv2.imwrite(save_name, img_stdlabel)

            save_name = join(side_DIR_INV, title) + ".jpg"
            ret, sideImg = cv2.threshold(sideImg, 128, 255, cv2.THRESH_BINARY_INV)
            cv2.imwrite(save_name, sideImg)

            imgShow = np.empty((size_ori[0], size_ori[1], 3), np.uint8)
            imgShow[:, :, 0] = img_ori
            imgShow[:, :, 1] = img_ori
            imgShow[:, :, 2] = np.maximum(img_ori, boxImg) * 0.2 + img_ori * 0.8
            save_name = join(label_DIR, title) + "_mark.png"
            # cv2.imwrite(save_name, imgShow)
            imgShow = cv2.resize(imgShow,(int(size_ori[0]/1.5), int(size_ori[1]/1.5)))
            imgShow = cv2.putText(imgShow, title, (50, 150), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
            cv2.imshow("imgShow", imgShow)
            cv2.waitKey(1)
            save_name = join(img_DIR, "%s.jpg"%title)
            cv2.imwrite(save_name, img_ori_bgr)

            # imgShow = np.empty((size_ori[0], size_ori[1], 3), np.uint8)
            # imgShow[:, :, 0] = img_ori
            # imgShow[:, :, 1] = np.maximum(img_ori, sideImg) * 0.4 + img_ori * 0.6
            # imgShow[:, :, 2] = img_ori
            # save_name = join(side_DIR, title) + "_mark.png"
            # cv2.imwrite(save_name, imgShow)

            sss = np.zeros((size_ori[0], size_ori[1], 1), np.uint8)

            image = cv2.add(img_ori, np.zeros(np.shape(img_ori), dtype=np.uint8), mask=boxImg)
            save_name = join(ROI_DIR, title) + "_mark2.png"
            cv2.imwrite(save_name, image)





if __name__ == "__main__":
    # labelme_2_coco()
    labelme_2_mask_rcnn_ballon()